from dotenv import load_dotenv, find_dotenv
from langchain_community.chat_models import ChatZhipuAI
from langchain_core.messages import  HumanMessage
from langchain_core.callbacks.streaming_stdout import (
    StreamingStdOutCallbackHandler
)

_ = load_dotenv(find_dotenv())

handlers = [StreamingStdOutCallbackHandler()]
messages = [
    HumanMessage(content="请写一首关于春天的诗，并需要有标题。"),
]

# streaming_chat = ChatZhipuAI(
#     model="glm-4-plus",
#     temperature=0.9,
#     streaming=True,
#     callbacks=handlers,
# )

# streaming_chat.invoke(messages)

streaming_chat = ChatZhipuAI(
    model="glm-4-plus",
    temperature=0.9,
    streaming=True,
)

streaming_chat.invoke(messages, config={"callbacks": handlers})
